# 导包
from tensorflow.keras import Model
from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, Dense

'''
卷积操作tf.keras.layers.Conv2D
用到的参数：
- input_shape: 输入数据的维度大小,如3通道64*64的输入图片，要设置为（64，64，3），
参数input_shape只需要在第一个网络层进行设置。
- filters: 输入整数，卷积核个数（等于卷积后输出的通道数）
- kernel_size: 卷积核的大小，通常为一个正整数或一个列表。
    ex： kernel_size = 3,kernel_size(3,3),都表示卷积核宽高方向的大小为3.
- strides： 卷积核的滑动步长，通常为一个正整数或者一个列表，列表可以包含2个或4个元素，默认为（1， 1）
    ex： strides=2,strides=(2, 2),都表示宽高方向的滑动步长为2.
- padding: 补零的设置，可选项为“Valid”，“Same”。
    -如果padding="Valid",表示不做补零操作；
    -如果padding="Same",程序会自动补零，使得输出feature_map的宽高等于ceil（输入图片宽高/strides）
    ceil表示向上取整，当strides=1时， 输入输出的宽高是一致。
- activation: 激活函数
    可选项：
    -'sigmoid': sigmoid激活函数
    -'tanh': tanh激活函数
    -'relu': relu激活函数
    -'elu'或者tf.keras.activations.elu(alpha=1.0): elu激活函数
    -‘selu’： selu激活函数
    -'swish': swish激活函数（tf2.2版本以上才有）
    -'softmax': softmax激活函数
    -'input_shape': 如果是第一层卷积，需要设置输入图片的大小（height, width, channel）, 
                    如input_shape(128,128,3)
    -'name': 输入字符串，给该层设置一个名称。
    
池化操作tf.keras.layers.Maxpool2D
用到的参数：
- pool_size: 池化卷积核的大小，具体设置tf.keras.layers.Conv2D的kernel_size一样。
- strides: 池化的滑动步长，具体设置tf.keras.layers.Conv2D的strides一样。
- name: 输入字符串，给该层设置一个名称。

全连接操作tf.keras.layers.Dense
用到的参数：
- units: 输入整数，全连接层神经元个数
- activation:  激活函数，分类网络的输出层一般用‘softmax’激活函数。
- name: 输入字符串，给该层设置一个名称。

展平层操作tf.keras.layers.Flatten
举例说明：
[[1, 2, 3]],
[[4, 5, 6]],     --->   [1, 2, 3, 4, 5, 6, 7, 8, 9]
[[7, 8, 9]]
'''

# 定义一个子类来搭建模型
class ConvModel(Model):
    def __init__(self):
        # 父类初始化
        super(ConvModel,self).__init__()

        # 卷积层con_1_1
        self.conv_1_1 = Conv2D(input_shape=(64,64,3), filters=64, kernel_size=3,
                               padding='same', activation='relu', name='conv_1_1')

        # 卷积层conv_1_2
        self.conv_1_2 = Conv2D(filters=64, kernel_size=3, padding='same',
                               activation='relu', name='conv_1_2')

        # 池化层max_pool_1
        self.max_pool_1 = MaxPooling2D(pool_size=2, name='max_pool_1')

        # 卷积层conv_2_1
        self.conv_2_1 = Conv2D(filters=128, kernel_size=3, padding='same',
                               activation='relu', name='conv_2_1')

        # 卷积层onv_2_2
        self.conv_2_2 = Conv2D(filters=128, kernel_size=3, padding='same',
                                activation='relu', name='conv_2_2')

        # 池化层max_pool_2
        self.max_pool_2 = MaxPooling2D(pool_size=2, name='max_pool_2')

        #  展平层flatten
        self.flatten = Flatten(name='flatten')

        # 全连接层
        self.dense = Dense(units=10, activation='softmax', name='logit')

    def call(self, x):
        x = self.conv_1_1(x)
        x = self.conv_1_2(x)
        x = self.max_pool_1(x)
        x = self.conv_2_1(x)
        x = self.conv_2_2(x)
        x = self.max_pool_2(x)
        x = self.flatten(x)
        x = self.dense(x)
        return x

